/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: TbeUnsortedSegmentSum的实现文件
 * Author: zhangjian
 * Create: 2023-5-17
 */

#include "tbe_unsorted_segment_sum_aicore.h"
#include <fstream>
#include <dlfcn.h>
#include <unistd.h>
#include <nlohmann/json.hpp>
#include <cstring>
#include <hccl/hccl_types.h>
#include <string>
#include <sstream>
#include "tbe_vector_reduce.h"
#include "runtime/rt.h"
#include "op_json_info.h"
#include "op_tiling.h"
#include "vector_tiling.h"
#include "adapter_rts.h"
#include "mmpa_api.h"
#include "unsorted_segment_sum_tiling.h"
#include "unsorted_segment_sum_tiling_rt2.h"
#include "tbe_gatherv2_aicore.h"

#ifndef HCCD
#include "platform/platform_info.h"
#include "exe_graph/runtime/tiling_context.h"
#include "exe_graph/runtime/tiling_parse_context.h"
#include "exe_graph/lowering/tiling_context_builder.h"
#include "exe_graph/lowering/tiling_parse_context_builder.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/operator_factory.h"
#include "graph/operator_reg.h"
#include "graph/operator.h"
#include "vector_tiling_rt2.h"
#include "graph/ge_tensor.h"
#include "graph/op_desc.h"
#include "graph/utils/graph_utils.h"
#endif

using namespace std;

template <typename T, typename std::enable_if<(!std::is_array<T>::value), int32_t>::type = 0>
static void *CreateCompileInfo()
{
    auto p = new T();
    return p;
}
template <typename T> static void DeleteCompileInfo(void * const obj)
{
    return;
}

namespace TbeReduce {

constexpr u32 MEMSET_UNSORTED_SEGMENT_SUM_TILING_DATA_MAX_SZIE = 1000;
constexpr u32 MEMSET_OVERFLOW_SZIE = 32;
constexpr u32 MEMSET_NUM_SEGMENTS_SZIE = 8;
constexpr u32 KMAX_TILING_DATA_SIZE = 16UL * 1024UL;
constexpr u32 KMAX_WORKSPACE_COUNT = 16;
fe::PlatFormInfos g_ussPlatFormInfos;


TbeUnsortedSegmentSum::TbeUnsortedSegmentSum()
{
}

TbeUnsortedSegmentSum::~TbeUnsortedSegmentSum()
{
    if (numSegments_ != nullptr) {
        hrtFree(numSegments_);
        numSegments_ = nullptr;
    }
    if (overFlowPtr_ != nullptr) {
        hrtFree(overFlowPtr_);
        overFlowPtr_ = nullptr;
    }
}

HcclResult TbeUnsortedSegmentSum::Init() // 初始化
{
    CHK_PRT_RET(isInit_, HCCL_DEBUG("TbeUnsortedSegmentSum had been initialized"), HCCL_SUCCESS);
    CHK_RET(hrtGetDeviceType(deviceType_));
    switch (deviceType_) {
        case DevType::DEV_TYPE_910:
            InitOpInfoMap910A();
            GetCoreNum();
            break;
        case DevType::DEV_TYPE_310P3:
            InitOpInfoMap310P3();
            break;
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            InitOpInfoMap910B();
            GetCoreNum();
            break;
        default:
            HCCL_ERROR("[TbeUnsortedSegmentSum][Init]devType[%u] is not supported", deviceType_);
            return HCCL_E_NOT_SUPPORT;
    }
    tilingDataHostPtr_ = new(nothrow) char[MEMSET_UNSORTED_SEGMENT_SUM_TILING_DATA_MAX_SZIE]; // 使用内存申请
    CHK_PTR_NULL(tilingDataHostPtr_);
    CHK_RET(LoadOpBinary()); // 加载二进制文件
    isInit_ = true;
    return HCCL_SUCCESS;
}

HcclResult TbeUnsortedSegmentSum::GetTilingRunInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo,
    u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo)
{
    TeOpTensor inputTensor;
    inputTensor.dtype = "float32";
    inputTensor.shape = { static_cast<s64>(inputCount), valueDim };
    inputTensor.oriShape = { static_cast<s64>(inputCount), valueDim };
    inputTensor.format = "ND";
    inputTensor.oriFormat = "ND";

    TeOpTensorArg inputTilingArg;
    inputTilingArg.tensor.push_back(inputTensor);
    inputTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpTensor segmentidsTensor;

    if (dataType == HCCL_DATA_TYPE_INT32) {
        segmentidsTensor.dtype = "int32";
    } else if (dataType == HCCL_DATA_TYPE_INT64) {
        segmentidsTensor.dtype = "int64";
    } else {
        HCCL_ERROR("datatype[%s] is not int32 or int64", GetDataTypeEnumStr(dataType).c_str());
        return HCCL_E_PARA;
    }

    segmentidsTensor.shape = { static_cast<s64>(inputCount) };
    segmentidsTensor.oriShape = { static_cast<s64>(inputCount) };
    segmentidsTensor.format = "ND";
    segmentidsTensor.oriFormat = "ND";

    TeOpTensorArg segmentidsTilingArg;
    segmentidsTilingArg.tensor.push_back(segmentidsTensor);
    segmentidsTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpTensor outputTensor;
    outputTensor.dtype = "float32";
    outputTensor.shape = { static_cast<s64>(outputCount), valueDim };
    outputTensor.oriShape = { static_cast<s64>(outputCount), valueDim };
    outputTensor.format = "ND";
    outputTensor.oriFormat = "ND";

    TeOpTensorArg outputTilingArg;
    outputTilingArg.tensor.push_back(outputTensor);
    outputTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpParas paras;
    paras.inputs.push_back(inputTilingArg);
    paras.inputs.push_back(segmentidsTilingArg);
    paras.outputs.push_back(outputTilingArg);
    paras.attrs.push_back("0");

    CompileParas compileInfo;
    compileInfo.is_tik = opTilingInfo["compileInfo"]["is_tik"];
    compileInfo.core_num = opTilingInfo["compileInfo"]["vars"]["core_num"];
    compileInfo.impl_mode = opTilingInfo["compileInfo"]["vars"]["impl_mode"];
    compileInfo.ub_size = opTilingInfo["compileInfo"]["vars"]["ub_size"];
    compileInfo.ub_tensor_num = opTilingInfo["compileInfo"]["vars"]["ub_tensor_num"];

    string opName;

    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName)); // 获取描述文件中数据
    bool isTiling = SegmentTIKTiling(opName, paras, compileInfo, runInfo);
    CHK_PRT_RET(!isTiling, HCCL_ERROR("[Get][TilingRunInfo]unsorted segment sum tiling failed"), HCCL_E_CCE);

    return HCCL_SUCCESS;
}

#ifndef HCCD
// rt2.0
HcclResult TbeUnsortedSegmentSum::GetTilingRunInfo(nlohmann::json &opDescInfo,
    u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo)
{
    ge::OpDescPtr op_desc = nullptr;
    EXECEPTION_CATCH(op_desc = make_shared<ge::OpDesc>("tmp_op", "unsorted_segment_sum"), return HCCL_E_PTR);
    std::vector<int64_t> dimsIn = { static_cast<s64>(inputCount), valueDim };
    std::vector<int64_t> dimsSeg = { static_cast<s64>(inputCount), };
    std::vector<int64_t> dimsOut = { static_cast<s64>(outputCount), valueDim};
    ge::GeTensorDesc tensor_descIn(ge::GeShape(dimsIn), ge::FORMAT_ND, ge::DataType::DT_FLOAT);
    ge::GeTensorDesc tensor_descSeg(ge::GeShape(dimsSeg), ge::FORMAT_ND, ge::DataType::DT_INT32);
    ge::GeTensorDesc tensor_descOut(ge::GeShape(dimsOut), ge::FORMAT_ND, ge::DataType::DT_FLOAT);

    op_desc->AddInputDesc("x", tensor_descIn);
    op_desc->AddInputDesc("segment_ids", tensor_descSeg);
    op_desc->AddOutputDesc("y", tensor_descOut);

    ge::ComputeGraphPtr graph = make_shared<ge::ComputeGraph>("tmp");
    ge::NodePtr node = graph->AddNode(op_desc);
    const size_t INPUT_SIZE = 2;
    CHK_RET(AddDataNodeForOriginNode(graph, node, INPUT_SIZE));
    ge::Operator op = ge::OpDescUtils::CreateOperatorFromNode(node);

    // 解析
    std::string compileInfoJson = opDescInfo.dump();
    HCCL_INFO("compileInfoJson[%s]", compileInfoJson.c_str());

    auto tilingParseContext =
        gert::TilingParseContextBuilder()
            .CompileJson(compileInfoJson.c_str())
            .PlatformInfo(reinterpret_cast<void *>(const_cast<fe::PlatFormInfos *>
                (&g_ussPlatFormInfos)))
            .CompileInfoCreatorFunc(CreateCompileInfo<optilingUss::TilingPrepareForUnsortedSegmentSumCompileInfo>)
            .CompileInfoDeleterFunc(DeleteCompileInfo<optilingUss::TilingPrepareForUnsortedSegmentSumCompileInfo>)
            .Build(op);

    auto parseContext = reinterpret_cast<gert::TilingParseContext *>(tilingParseContext.GetKernelContext());
    CHK_PTR_NULL(parseContext);

    ge::graphStatus gRet = optilingUss::TilingPrepare4SegmentSumComm(parseContext);
    CHK_PRT_RET(gRet != ge::GRAPH_SUCCESS, HCCL_ERROR("[GetTilingRunInfo] Auto TilingParser failed"), HCCL_E_PARA);

    auto compiledInfoUss = parseContext->GetCompiledInfo<optilingUss::TilingPrepareForUnsortedSegmentSumCompileInfo>();
    CHK_PTR_NULL(compiledInfoUss);
    auto compiledInfo = compiledInfoUss->dsl_compile_info;
    CHK_PTR_NULL(compiledInfo);

    // 动态获取core的数量
    compiledInfo->soc_info.core_num = g_ussPlatFormInfos.GetCoreNum();
    if (compiledInfo->soc_info.core_num <= 0) {
        HCCL_ERROR("Fail to get core num from GE, core num get is %d.", compiledInfo->soc_info.core_num);
        return HCCL_E_PARA;
    }

    HCCL_INFO("Core num is %d.", compiledInfo->soc_info.core_num);

    // 构造tiling入参，进行tiling操作
    uint32_t maxSize = -1;
    if (GetOpProperty(opDescInfo, "opParaSize", maxSize) != HCCL_SUCCESS) {
        HCCL_INFO("No max tiling size in opdesc.");
        maxSize = KMAX_TILING_DATA_SIZE;
    }
    const auto alignedMaxSize = RoundUp(static_cast<uint64_t>(maxSize), sizeof(uintptr_t));
    const auto tilingData = gert::TilingData::CreateCap(alignedMaxSize);
    const auto workspaceSize = gert::ContinuousVector::Create<size_t>(KMAX_WORKSPACE_COUNT);
    auto context_builder = gert::TilingContextBuilder();
    const gert::KernelContextHolder tilingContext =
        context_builder
            .CompileInfo(reinterpret_cast<void *>(compiledInfoUss))
            .PlatformInfo(reinterpret_cast<void *>(const_cast<fe::PlatFormInfos *>(&g_ussPlatFormInfos)))
            .TilingData(tilingData.get())
            .Workspace(reinterpret_cast<gert::ContinuousVector *>(workspaceSize.get()))
            .Build(op);
    gert::TilingContext *context = reinterpret_cast<gert::TilingContext *>(tilingContext.context_);
    CHK_PTR_NULL(context);

    auto rawTilingData = context->GetRawTilingData();
    if (rawTilingData != nullptr) {
        rawTilingData->SetDataSize(0);
    }

    gRet = optilingUss::Tiling4UnsortedSegmentSum(context);
    CHK_PRT_RET(gRet != ge::GRAPH_SUCCESS, HCCL_ERROR("[GetTilingRunInfo] Auto Tiling failed"), HCCL_E_PARA);
    CHK_RET(ConvertToOpRunInfo(context, runInfo));
    return HCCL_SUCCESS;
}
#endif

HcclResult TbeUnsortedSegmentSum::Run(const void *input, const void *segmentids,
    u64 inputCount, u64 outputCount, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst)
{
    CHK_PRT_RET(!isInit_, HCCL_ERROR("[Run][Reduce]TbeVectorReduce had not been initialized"), HCCL_E_UNAVAIL);
    CHK_PTR_NULL(input);
    CHK_PTR_NULL(segmentids);
    CHK_PTR_NULL(dst);
    CHK_PTR_NULL(stream);
    CHK_PRT_RET(dataType >= HCCL_DATA_TYPE_RESERVED, HCCL_ERROR("[Run][Reduce]unsupported data type[%s]",
        GetDataTypeEnumStr(dataType).c_str()), HCCL_E_PARA);

    HCCL_DEBUG("inputCount[%llu], outputCount[%llu], valueDim[%u]", inputCount, outputCount, valueDim);

    // 获取op关联的描述信息和tiling信息
    nlohmann::json opDescInfo;
    nlohmann::json opTilingInfo;
    CHK_RET(GetOpInfo(opDescInfo, opTilingInfo));
    HCCL_INFO("opDescInfo[%s]", opDescInfo.dump().c_str());

    // 获取tiling data
    OpRunInfo runInfo;
    runInfo.blockDim = 0;
    runInfo.workspaces = { 0 };
    runInfo.clearAtomic = 0;
    // 申请numSegments内存并拷贝outputCount
    CHK_RET(hrtMalloc(&numSegments_, sizeof(s64)));
    CHK_RET(hrtMemcpy(numSegments_, sizeof(s64), &outputCount, sizeof(s64),
        HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE));
    if (SupportRt2()) {
#ifndef HCCD
        // 910B & 910_93
        CHK_RET(GetTilingRunInfo(opDescInfo, inputCount, outputCount, valueDim, dataType, runInfo));
#else
        HCCL_ERROR("[TbeUnsortedSegmentSum]Does not support this interface.");
        return HCCL_E_NOT_SUPPORT;
#endif
    } else {
        // 910A & 310P
        CHK_RET(GetTilingRunInfo(opDescInfo, opTilingInfo, inputCount, outputCount, valueDim, dataType, runInfo));
        // 申请overflow内存
        CHK_RET(hrtMalloc(&overFlowPtr_, MEMSET_OVERFLOW_SZIE));
    }

    // 将tilingdata拷贝device侧
    unique_lock<mutex> tilingDataLock(tilingDataMutex_);
    void *tilingDataDevPtr = nullptr;
    CHK_RET(GetTilingDataDevMem(runInfo, &tilingDataDevPtr, MEMSET_UNSORTED_SEGMENT_SUM_TILING_DATA_MAX_SZIE));
    tilingDataLock.unlock();

    // 构造桩函数
    string opName;
    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName));

    s8 *stubFunc = nullptr;
    string kernelName;
    CHK_RET(GetOpProperty(opDescInfo, "kernelName", kernelName));
    if (SupportRt2()) {
#ifndef HCCD
        kernelName += '_';
        kernelName += to_string(runInfo.tilingKey);
#endif
    }

    unique_lock<mutex> stubFuncLock(stubFuncMutex_);
    CHK_RET(GetstubFunc(opName, stubFunc, kernelName));
    stubFuncLock.unlock();

    // 执行OP
    CHK_RET(
        ExecuteKernelLaunch(input, segmentids, numSegments_, dst, stream, reinterpret_cast<void *>(stubFunc), runInfo,
            tilingDataDevPtr, overFlowPtr_));
    return HCCL_SUCCESS;
}

HcclResult TbeUnsortedSegmentSum::ExecuteKernelLaunch(const void *input, const void *segmentids,
    const void *num_segments, const void *dst, void *stream, const void *stubFunc, OpRunInfo &runInfo,
    void *tilingDataDevPtr, const void *overFlowPtr)
{
    vector<const void *> deviceAddrs;
    deviceAddrs.push_back(input);
    deviceAddrs.push_back(segmentids);
    deviceAddrs.push_back(num_segments);
    deviceAddrs.push_back(dst);
    deviceAddrs.push_back(tilingDataDevPtr);
    if (!SupportRt2()) {
        deviceAddrs.push_back(overFlowPtr);
    }

    HcclResult ret = hrtKernelLaunch(stubFunc, runInfo.blockDim, deviceAddrs.data(),
        deviceAddrs.size() * sizeof(void *), nullptr, stream);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Execute][KernelLaunch]execute kernel launch failed"), HCCL_E_RUNTIME);
    return HCCL_SUCCESS;
}

HcclResult TbeUnsortedSegmentSum::GetopInfoIndex(string &opInfoIndex)
{
    // 芯片类型适配
    switch (deviceType_) {
        case DevType::DEV_TYPE_910:
            opInfoIndex += "te_unsortedsegmentsum_da0a5264250330a6c92ab5b622124e7d75359168a46ae6eab47256ad57367ba9_1";
            break;
        case DevType::DEV_TYPE_310P3:
            opInfoIndex += "te_unsortedsegmentsum_d0a8c7e68142042963e855ce6004fa68d52228200ba7cf5426e5ab47e86ae38c_1";
            break;
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            opInfoIndex += "uss_Ascend910B4_float32";
            break;
        default:
            HCCL_ERROR("deviceType_[%d] is not support", deviceType_);
            return HCCL_E_NOT_SUPPORT;
    }

    return HCCL_SUCCESS;
}

HcclResult TbeUnsortedSegmentSum::GetCoreNum()
{
#ifndef HCCD
    s32 deviceLogicId = -1;
    s64 coreNum = -1;
    HcclRtDeviceModuleType hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_DEVICE_MOUDLE_RESERVED;
    CHK_RET(hrtGetDevice(&deviceLogicId));

    if (deviceType_ == DevType::DEV_TYPE_910) {
        hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_MODULE_TYPE_AICORE;
    } else if (deviceType_ == DevType::DEV_TYPE_910B || deviceType_ == DevType::DEV_TYPE_910_93) {
        hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_MODULE_TYPE_VECTOR_CORE;
    } else {
        HCCL_ERROR("devType[%d] does not need tbe.", deviceType_);
        return HCCL_E_INTERNAL;
    }

    CHK_RET(hrtGetDeviceInfo(deviceLogicId, hcclModuleType, HcclRtDeviceInfoType::HCCL_INFO_TYPE_CORE_NUM, coreNum));
    HCCL_INFO("[GetCoreNum]coreNum[%lld]", coreNum);
    g_ussPlatFormInfos.SetCoreNum(static_cast<uint32_t>(coreNum));
#endif
    return HCCL_SUCCESS;
}

HcclResult TbeUnsortedSegmentSum::GetOpInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo)
{
    string opInfoIndex;
    CHK_RET(GetopInfoIndex(opInfoIndex));
    if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
        HCCL_ERROR("[Get][OpInfo]can't find tbereduce info,opInfoIndex[%s]", opInfoIndex.c_str());
        return HCCL_E_UNAVAIL;
    }
    opDescInfo = opInfoMap_[opInfoIndex];

    if (!SupportRt2()) {
        opInfoIndex += "_tiling_info";
        if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
            HCCL_ERROR("[Get][OpInfo]can't find tbereduce info,opInfoTilingIndex[%s]", opInfoIndex.c_str());
            return HCCL_E_UNAVAIL;
        }
        opTilingInfo = opInfoMap_[opInfoIndex];
    }

    return HCCL_SUCCESS;
}

void TbeUnsortedSegmentSum::InitOpInfoMap910A()
{
    opInfoMap_["te_unsortedsegmentsum_da0a5264250330a6c92ab5b622124e7d75359168a46ae6eab47256ad57367ba9_1"] =
        te_unsortedsegmentsum_da0a5264250330a6c92ab5b622124e7d75359168a46ae6eab47256ad57367ba9_1;
    opInfoMap_["te_unsortedsegmentsum_da0a5264250330a6c92ab5b622124e7d75359168a46ae6eab47256ad57367ba9_1_tiling_info"] =
        te_unsortedsegmentsum_da0a5264250330a6c92ab5b622124e7d75359168a46ae6eab47256ad57367ba9_1_tiling_info;
}

void TbeUnsortedSegmentSum::InitOpInfoMap310P3()
{
    opInfoMap_["te_unsortedsegmentsum_d0a8c7e68142042963e855ce6004fa68d52228200ba7cf5426e5ab47e86ae38c_1"] =
        te_unsortedsegmentsum_d0a8c7e68142042963e855ce6004fa68d52228200ba7cf5426e5ab47e86ae38c_1;
    opInfoMap_["te_unsortedsegmentsum_d0a8c7e68142042963e855ce6004fa68d52228200ba7cf5426e5ab47e86ae38c_1_tiling_info"] =
        te_unsortedsegmentsum_d0a8c7e68142042963e855ce6004fa68d52228200ba7cf5426e5ab47e86ae38c_1_tiling_info;
}

void TbeUnsortedSegmentSum::InitOpInfoMap910B()
{
    opInfoMap_["uss_Ascend910B4_float32"] = uss_Ascend910B4_float32;
}

} // namespace TbeReduce

#ifdef __cplusplus
extern "C" {
#endif
// 无业务验证使用
static TbeReduce::TbeUnsortedSegmentSum g_unsortedSegmentSum;
void TbeUnsortedSegmentSumInit()
{
    g_unsortedSegmentSum.Init();
}

void TbeUnsortedSegmentSumRun(const void *input, const void *segmentids,
    u64 inputCount, u64 outputCount, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst)
{
    g_unsortedSegmentSum.Run(input, segmentids, inputCount, outputCount, valueDim, dataType, stream, dst);
}
#ifdef __cplusplus
}
#endif
